package org.taj.adapter.server;

import com.alibaba.fastjson2.JSON;
import org.taj.adapter.llama.LLaMAClient;
import org.taj.adapter.llama.LLaMARequest;

import java.util.ArrayList;
import java.util.List;

public class Session {
    private LLaMARequest.Message prompt;
    private List<LLaMARequest.Message> history;
    private long created;
    private long promptN;
    private double promptElapsedSecond;
    private long predictedN;
    private double predictedElapsedSecond;

    public Session() {
        history = new ArrayList<>();
        prompt = new LLaMARequest.Message(0, "system", "");
        history.add(prompt);
        created = System.currentTimeMillis();
    }

    public void setPrompt(String prompt) {
        this.prompt.setContent(prompt);
    }

    public void addUserInput(String input) {
        LLaMARequest.Message message = new LLaMARequest.Message(System.currentTimeMillis(), "user", input);
        history.add(message);
    }

    public void addAssistantResponse(String response) {
        LLaMARequest.Message message = new LLaMARequest.Message(System.currentTimeMillis(), "assistant", response);
        history.add(message);
    }

    public void setTimings(long promptN, double promptElapsedSecond, long predictedN, double predictedElapsedSecond) {
        this.promptN += promptN;
        this.promptElapsedSecond += promptElapsedSecond;
        this.predictedN += predictedN;
        this.predictedElapsedSecond += predictedElapsedSecond;
    }

    public String toString() {
        LLaMARequest request = new LLaMARequest(history);
        return JSON.toJSONString(request);
    }


}
